/**
 * Copyright (C) 2011 Andrew C. Love (DNC) <dnc.app.sup@gmail.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.dnc.cloak.service;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.Serializable;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Collections;
import java.util.Date;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.apache.log4j.Logger;

import com.dnc.cloak.framework.context.Context;
import com.dnc.cloak.framework.persistence.gen.Keystore;
import com.dnc.cloak.framework.persistence.gen.KeystoreEntry;
import com.dnc.cloak.framework.security.encryption.AsymmetricKeyEncryptor;
import com.dnc.cloak.framework.security.encryption.EncryptionException;
import com.dnc.cloak.framework.security.encryption.EncryptionProperties;
import com.dnc.cloak.framework.security.encryption.EncryptionProvider;
import com.dnc.cloak.framework.security.encryption.EncryptionUtil;
import com.dnc.cloak.framework.security.encryption.KeyContext;
import com.dnc.cloak.framework.security.encryption.KeystoreBean;
import com.dnc.cloak.framework.security.encryption.SymmetricKeyEncryptor;
import com.dnc.cloak.framework.service.ServiceException;

public class EncryptionServiceHelper implements Serializable {
  
	private static final Logger logger = Logger.getLogger(EncryptionServiceHelper.class.getName());

	private static final Map<String,KeystoreBean> keystoreBeanMap = Collections.synchronizedMap(new HashMap<String,KeystoreBean>());
	
    private static SecretKey defaultSymmetricKey = null;
    private static final byte[] aesDefaultKeyBytes = { (byte) 0x74, (byte) 0x72, (byte) 0x65, (byte) 0x78, (byte) 0x6f, (byte) 0x6e, (byte) 0x65, (byte) 0x36, (byte) 0x74, (byte) 0x72, (byte) 0x65, (byte) 0x78, (byte) 0x6f, (byte) 0x6e, (byte) 0x65, (byte) 0x36 };


    public EncryptionServiceHelper() throws Exception {

    }

    public static synchronized void removeKeyStoreBean(String keystoreName) {
        keystoreBeanMap.remove(keystoreName);
    }
    
    public static synchronized void addKeystore(KeystoreBean keystoreBean) throws EncryptionException {
        keystoreBeanMap.put(keystoreBean.getKeystore().getKeystoreName(), keystoreBean);
    }
    
    public static synchronized Certificate getCertificate(KeystoreBean keystoreBean, String keyAlias) throws KeyStoreException {
        return keystoreBean.getJavaKeystore().getCertificate(keyAlias);
    }

    public static synchronized PublicKey getPublicKey(KeystoreBean keystoreBean, String keyAlias) throws KeyStoreException {
        try {
            return keystoreBean.getJavaKeystore().getCertificate(keyAlias).getPublicKey();
        }
        catch (KeyStoreException e) {
            logger.error("Failed to getPublicKey:" + keyAlias + " from keystore:" + keystoreBean.getKeystore().getKeystoreName(), e);
            throw new KeyStoreException("Failed to getPublicKey:" + keyAlias + " from keystore:" + keystoreBean.getKeystore().getKeystoreName(),e);
        }
    }

    public static synchronized PrivateKey getPrivateKey(KeystoreBean keystoreBean, String keyAlias) throws KeyStoreException, NoSuchAlgorithmException, UnrecoverableKeyException {
        KeystoreEntry keystoreEntry = keystoreBean.getKeystoreEntryMap().get(keyAlias);

        try {
            return (PrivateKey) keystoreBean.getJavaKeystore().getKey(keyAlias, keystoreEntry.getKeystoreEntryToken().toCharArray());
        }
        catch (KeyStoreException ex) {
            logger.error("getPrivateKey", ex);
            logger.error("keyAlias:" + keyAlias);
            logger.error("Failed to getPrivateKey:" + keyAlias + " from keystore:" + keystoreBean.getKeystore().getKeystoreName(), ex);
            throw new KeyStoreException("Failed to getPrivateKey:" + keyAlias + " from keystore:" + keystoreBean.getKeystore().getKeystoreName(),ex);
        }

    }
    
    public static synchronized SecretKey getSymmetricKey(KeystoreBean keystoreBean, String keyAlias) throws KeyStoreException, NoSuchAlgorithmException,UnrecoverableKeyException {
        // Default is AES, but this method may be enhanced later to provide
        // a default symmetric key per keystore
        KeystoreEntry keystoreEntry = keystoreBean.getKeystoreEntryMap().get(keyAlias);

        if (logger.isDebugEnabled())
            logger.debug("keystore=" + keystoreBean.getKeystore().getKeystoreName());

        try {
            EncryptionUtil.setupProvider(keystoreBean.getKeystore().getKeystoreJceProvider());
            defaultSymmetricKey = new SecretKeySpec(aesDefaultKeyBytes,keystoreEntry.getKeystoreEntrySymAlgorithm());
        } 
        catch (Exception ex) {
            logger.error("Exception occurred while attempting to get default symmetric Key.",ex);
            throw new KeyStoreException("Exception occurred while attempting to get default symmetricKey.",ex);
        }
        
        return defaultSymmetricKey;
        
    }
  
  
    public static synchronized SecretKey createSymmetricKey(KeystoreBean keystoreBean, String keyAlias) throws ServiceException, Exception {
      
        try {
            KeystoreEntry keystoreEntry = keystoreBean.getKeystoreEntryMap().get(keyAlias);
            
            EncryptionUtil.setupProvider(keystoreBean.getKeystore().getKeystoreJceProvider());

            String algorithm = keystoreEntry.getKeystoreEntrySymAlgorithm();
            
            KeyGenerator keyGenerator = KeyGenerator.getInstance(algorithm);
            int keysize = keystoreEntry.getKeystoreEntrySymKeyInitSize();
            keyGenerator.init(keysize);
            SecretKey secretKey = keyGenerator.generateKey();

            if (logger.isDebugEnabled()) {
                logger.debug("secret key:" + secretKey.toString());
            }

            return secretKey;
        } 
        catch (Exception ex) {
            throw new KeyStoreException("Exception occurred while attempting to createSymmetricKey.",ex);
        }
        
    }   
    
    public static synchronized KeystoreBean getKeyStoreBean(String keyStoreName) {
        return keystoreBeanMap.get(keyStoreName);
    }        
 
    public static synchronized KeystoreBean createKeystoreBean(Context context, Keystore keystore, List<KeystoreEntry> keystoreEntryList) throws EncryptionException {
        
        ByteArrayInputStream keyStoreIS = new ByteArrayInputStream(keystore.getKeystoreContent());
        KeystoreBean keystoreBean;
        try {
            KeyStore javaKeyStore = EncryptionUtil.loadKeyStorefromStream(keyStoreIS, keystore.getKeystoreJceProvider(), keystore.getKeystoreType(), keystore.getKeystoreToken());
            keystoreBean = new KeystoreBean(keystore, keystoreEntryList,javaKeyStore);
        }
        catch (EncryptionException e) {
          logger.error("Failed to load keystore" + ":name:" + keystore.getKeystoreName() + ":keyStoreType:" + keystore.getKeystoreType() + " " + e.toString(), e);
          logger.error("Failed to load keystore" + ":keyStoreBytesLength:" + ((keystore.getKeystoreContent() == null) ? "null" : keystore.getKeystoreContent().length + ""), e);
          logger.error("Failed to load keystore from stream" + ":keyStoreIS.available():" + keyStoreIS.available(), e);
          throw e;
        }
            
        return keystoreBean;

    }    

    public static synchronized KeystoreBean createKeystoreBean(Context context, KeyStore javaKeystore, KeyContext keyContext) throws Exception, KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException {
        
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        javaKeystore.store(baos, keyContext.getToken().toCharArray());

        Keystore keystore = new Keystore();
        keystore.setInsert(true);
        keystore.setKeystoreContent(baos.toByteArray());
        keystore.setKeystoreJceProvider(keyContext.getProviderClassName());
        keystore.setKeystoreLocation(keyContext.getLocation());
        keystore.setKeystoreName(keyContext.getKeystoreName());
        keystore.setKeystoreToken(keyContext.getToken());
        keystore.setKeystoreType(keyContext.getKeystoreType());
        
        Enumeration<String> aliases = javaKeystore.aliases();
        List<KeystoreEntry> keystoreEntries = new ArrayList<KeystoreEntry>();
        
        while(aliases!=null && aliases.hasMoreElements()) {
            String alias = (String)aliases.nextElement();

            X509Certificate cert = (X509Certificate)javaKeystore.getCertificate(alias);
            
            
            KeystoreEntry keystoreEntry = new KeystoreEntry();
            keystoreEntry.setInsert(true);
            keystoreEntries.add(keystoreEntry);
            keystoreEntry.setKeystoreName(keyContext.getKeystoreName());
            keystoreEntry.setKeystoreEntryAlias(alias);
            
            keystoreEntry.setKeystoreEntryDistinguishedName(cert.getIssuerDN().getName());
            keystoreEntry.setKeystoreEntryBeginDate(cert.getNotBefore());
            keystoreEntry.setKeystoreEntryExpireDate(cert.getNotAfter());
            keystoreEntry.setKeystoreEntryNewOnExpire(false);            
            long certDaysValid = (cert.getNotAfter().getTime() - cert.getNotBefore().getTime()) / (1000 * 60 * 60 * 24);
            keystoreEntry.setKeystoreEntryDaysValid((int)certDaysValid);
            keystoreEntry.setKeystoreEntryPkiAlgorithm(EncryptionProperties.DEFAULT_PKI_ALGORITHM);
            keystoreEntry.setKeystoreEntryPkiKeyInitSize(EncryptionProperties.DEFAULT_PKI_KEY_INIT_SIZE);
            keystoreEntry.setKeystoreEntryPkiTransformation(EncryptionProperties.DEFAULT_PKI_TRANSFORMATION);
            keystoreEntry.setKeystoreEntrySignatureAlgorithm(EncryptionProperties.DEFAULT_SIGNATURE_ALGORITHM);
            keystoreEntry.setKeystoreEntrySymAlgorithm(EncryptionProperties.DEFAULT_SYM_ALGORITHM);
            keystoreEntry.setKeystoreEntrySymTransformation(EncryptionProperties.SYM_TRANSFORMATION);
            keystoreEntry.setKeystoreEntrySymKeyInitSize(EncryptionProperties.DEFAULT_SYM_KEY_INIT_SIZE);        
            keystoreEntry.setKeystoreEntryToken(keyContext.getEntryToken());
        }
        
        KeystoreBean keystoreBean = new KeystoreBean(keystore, keystoreEntries, javaKeystore);
        
        return keystoreBean;        
    }

    public static synchronized KeystoreBean createKeyStore(Context context, String keystoreName, String alias) throws Exception, KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException {
        // Generate a random password for the keystore
        String keystoreGeneratedPassword = EncryptionUtil.generatePassword(16);
        String keystoreEntryGeneratedPassword = EncryptionUtil.generatePassword(16);

        String providerClassName = EncryptionProperties.DEFAULT_JCE_PROVIDER;
        String pkiAlgorithm = EncryptionProperties.DEFAULT_PKI_ALGORITHM;
        String pkiTransformation = EncryptionProperties.DEFAULT_PKI_TRANSFORMATION;
        int pkiKeySize = EncryptionProperties.DEFAULT_PKI_KEY_INIT_SIZE;        
        String symAlgorithm = EncryptionProperties.DEFAULT_SYM_ALGORITHM;
        String symTransformation = EncryptionProperties.DEFAULT_SYM_TRANSFORMATION;
        int symKeySize = EncryptionProperties.DEFAULT_SYM_KEY_INIT_SIZE;
        String signatureAlgorithm = EncryptionProperties.DEFAULT_SIGNATURE_ALGORITHM;
        String x509DistinguishedName = EncryptionProperties.DEFAULT_DISTINGUISHED_NAME;
        String keyStoreType = EncryptionProperties.DEFAULT_KEYSTORE_TYPE;
        String location = EncryptionProperties.DEFAULT_KEYSTORE;
        int certDaysValid = EncryptionProperties.DEFAULT_CERT_DAYS_VALID;

        // Creating a keystore is provider dependent, and complicated so defer
        // to helper.
        KeyStore generatedKeyStore = EncryptionUtil.createKeyStore(providerClassName, keyStoreType, x509DistinguishedName, signatureAlgorithm, alias, pkiAlgorithm, certDaysValid, pkiKeySize, keystoreGeneratedPassword,keystoreEntryGeneratedPassword);

        // Serialize the keystore securely
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        generatedKeyStore.store(baos, keystoreGeneratedPassword.toCharArray());

        Keystore keystore = new Keystore();
        keystore.setInsert(true);
        keystore.setKeystoreContent(baos.toByteArray());
        keystore.setKeystoreJceProvider(providerClassName);
        keystore.setKeystoreLocation(location);
        keystore.setKeystoreName(keystoreName);
        keystore.setKeystoreToken(keystoreGeneratedPassword);
        keystore.setKeystoreType(keyStoreType);        
        
        KeystoreEntry keystoreEntry = new KeystoreEntry();
        List<KeystoreEntry> keystoreEntries = new ArrayList<KeystoreEntry>();
        keystoreEntries.add(keystoreEntry);
        
        keystoreEntry.setInsert(true);
        keystoreEntry.setKeystoreName(keystoreName);
        keystoreEntry.setKeystoreEntryAlias(alias);
        keystoreEntry.setKeystoreEntryBeginDate(new Date());
        keystoreEntry.setKeystoreEntryDaysValid(certDaysValid);
        keystoreEntry.setKeystoreEntryDistinguishedName(x509DistinguishedName);
        Calendar cal = Calendar.getInstance();
        cal.setTime(new Date());
        cal.add(Calendar.DATE, certDaysValid);
        keystoreEntry.setKeystoreEntryExpireDate(cal.getTime());
        keystoreEntry.setKeystoreEntryNewOnExpire(true);
        keystoreEntry.setKeystoreEntryPkiAlgorithm(pkiAlgorithm);
        keystoreEntry.setKeystoreEntryPkiKeyInitSize(pkiKeySize);
        keystoreEntry.setKeystoreEntryPkiTransformation(pkiTransformation);
        keystoreEntry.setKeystoreEntrySignatureAlgorithm(signatureAlgorithm);
        keystoreEntry.setKeystoreEntrySymAlgorithm(symAlgorithm);
        keystoreEntry.setKeystoreEntrySymTransformation(symTransformation);
        keystoreEntry.setKeystoreEntrySymKeyInitSize(symKeySize);        
        keystoreEntry.setKeystoreEntryToken(keystoreEntryGeneratedPassword);
        
        KeystoreBean keystoreBean = new KeystoreBean(keystore, keystoreEntries, generatedKeyStore);
        
        return keystoreBean;

    }

      
    public static synchronized EncryptionProvider getEncryptionProvider(Context context, KeyContext keyContext, KeystoreBean keystoreBean, String keyAlias){
        KeystoreEntry keystoreEntry = keystoreBean.getKeystoreEntryMap().get(keyAlias);
        
        if (keyContext.isAsymmetricEncryptionMode()){
            return new AsymmetricKeyEncryptor(
                                                keystoreBean.getKeystore().getKeystoreName(),
                                                keystoreEntry.getKeystoreEntryPkiTransformation(),
                                                keystoreEntry.getKeystoreEntryPkiAlgorithm(),
                                                keystoreBean.getKeystore().getKeystoreJceProvider()
                                             );
        }
        else {
            return new SymmetricKeyEncryptor(
                    keystoreBean.getKeystore().getKeystoreName(),
                    keystoreEntry.getKeystoreEntrySymTransformation(),
                    keystoreEntry.getKeystoreEntrySymAlgorithm(),
                    keystoreBean.getKeystore().getKeystoreJceProvider()
                 );               
        }
        
    }        
    
}